import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import os
from torchvision import datasets
from torch.utils.data import DataLoader
from torch.autograd import Variable
from torchvision.utils import save_image
import torchvision.transforms as transforms
from torch.autograd import Variable
import numpy as np

#创建参数解析器对象
parser = argparse.ArgumentParser(description="超参数:")
parser.add_argument("--epochs", type=int, default=150, help="numbers of epochs of training")
parser.add_argument("--batch_size", type=int, default=32, help='size of the batches')
parser.add_argument("--lr", type=float, default=0.0002, help="Adam: learning rate") 
parser.add_argument("--b1", type=float, default=0.5, help="Adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="Adam:decay of first order momentum of gradieng")
parser.add_argument("--workers", type=int, default=8, help="number of cpu threads to use during batch generation")
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space") #生成器输入的潜在空间维度
parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="interval between image sampling")
opt = parser.parse_args()
print(opt)

cuda = True if torch.cuda.is_available() else False
#初始化神经网络权重
def weight_init_normal(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1: #卷积层，使用正太分布权重
        torch.nn.init.normal_(m.weight.data,0.0,0.02)
    elif classname.find("BatchNorm2d") != -1:
        torch.nn.init.normal_(m.weight.data,1.0,0.02) #批归一化层，正太分布初始权重，0初始偏置
        torch.nn.init.constant_(m.bias.data,0.0)
        
        
#生成器
class Generator(nn.Module):
    def __init__(self):
        super(Generator,self).__init__()
        
        self.init_size = opt.img_size // 4
        self.l1 = nn.Sequential(nn.Linear(opt.latent_dim,128*self.init_size**2))
        
        self.conv_blocks = nn.Sequential(
            nn.BatchNorm2d(128),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(128,128,kernel_size=3,stride=1,padding=1),
            nn.BatchNorm2d(128,0.8),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(128,64,kernel_size=3,stride=1,padding=1),
            nn.BatchNorm2d(64,0.8),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(64,opt.channels,kernel_size=3,stride=1,padding=1),
            nn.Tanh(),
        )
        
    def forward(self,z):
        out = self.l1(z)
        out = out.view(out.shape[0],128,self.init_size,self.init_size)
        img = self.conv_blocks(out)
        return img
    
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator,self).__init__()
        
        def discriminator(in_filters,out_filters,bn=True):
            block = [nn.Conv2d(in_filters,out_filters,kernel_size=3,stride=2,padding=1),
                     nn.LeakyReLU(0.2,inplace=True),
                     nn.Dropout(0.25)]
            if bn:
                block.append(nn.BatchNorm2d(out_filters,0.8))
            return block

        self.model = nn.Sequential(
            *discriminator(opt.channels,16,bn=False),
            *discriminator(16,32),
            *discriminator(32,64),
            *discriminator(64,128),
        )
        
        ds_size = opt.img_size // 2 ** 4
        self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2,1),
                                       nn.Sigmoid())
    def forward(self,img):
        out = self.model(img)
        out = out.view(out.shape[0],-1)
        validity = self.adv_layer(out)
        
        return validity
    
#定义损失函数
adversial_loss = torch.nn.BCELoss()

#初始化生成器和识别器
generator = Generator()
discriminator = Discriminator()

if cuda:
    generator.cuda()
    discriminator.cuda()
    adversial_loss.cuda()
    
#初始化权重
generator.apply(weight_init_normal)
discriminator.apply(weight_init_normal)

#从配置数据加载程序
os.makedirs("data/mnist",exist_ok=True)
dataloader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "data/mnist",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.Resize(opt.img_size),transforms.ToTensor(),transforms.Normalize([0.5],[0.5])]
        ),
    ),
    batch_size=opt.batch_size,
    shuffle=True,
)

#定义优化器
Optimizer_G = torch.optim.Adam(generator.parameters(),lr=opt.lr,betas=(opt.b1,opt.b2))
Optimizer_D = torch.optim.Adam(discriminator.parameters(),lr=opt.lr,betas=(opt.b1,opt.b2))

Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor

for epoch in range(opt.epochs):
    for i,(imgs,_) in enumerate(dataloader):
        
        valid = Variable(Tensor(imgs.shape[0],1).fill_(1.0),requires_grad=False)
        fake = Variable(Tensor(imgs.shape[0],1).fill_(0.0),requires_grad=False)
        
        real_imgs = Variable(imgs.type(Tensor))
        
        #训练生成器
        Optimizer_G.zero_grad()
        
        z = Variable(Tensor(np.random.normal(0,1,(imgs.shape[0],opt.latent_dim))))
        
        gen_imgs = generator(z)
        
        g_loss = adversial_loss(discriminator(gen_imgs),valid)
        
        g_loss.backward()
        Optimizer_G.step()
        
        
        #训练识别器
        Optimizer_D.zero_grad()
        
        real_loss = adversial_loss(discriminator(real_imgs),valid)
        fake_loss = adversial_loss(discriminator(gen_imgs.detach()),fake)
        d_loss = (real_loss + fake_loss ) / 2
        
        d_loss.backward()
        Optimizer_D.step()
        
        print(
            "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
            % (epoch, opt.epochs, i, len(dataloader), d_loss.item(), g_loss.item())
        )
        
        batches_done = epoch * len(dataloader) + i
        if batches_done % opt.sample_interval == 0:
            save_image(gen_imgs.data[:25], "images/%d.png" % batches_done, nrow=5, normalize=True)